#![feature(allocator_api)]

use std::{alloc::GlobalAlloc, sync::OnceLock};

use simple_alloc::{allocator::Allocator, memory_manager::MemoryManager, page::Pages};

use rand::Rng;

struct GlobalAllocator(OnceLock<Allocator>);

impl GlobalAllocator {
    // 64KiB
    const INITIAL_SIZE: usize = 64 * 1024;

    fn init() -> Allocator {
        let mem = Pages::new(MemoryManager::BINS_LEN + 1, Self::INITIAL_SIZE);
        Allocator::new(mem)
    }
}

unsafe impl GlobalAlloc for GlobalAllocator {
    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
        self.0.get_or_init(Self::init).alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        self.0.get_or_init(Self::init).dealloc(ptr, layout)
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
        self.0
            .get_or_init(Self::init)
            .realloc(ptr, layout, new_size)
    }

    unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 {
        self.0.get_or_init(Self::init).alloc_zeroed(layout)
    }
}

fn main() {
    let global_allocator: GlobalAllocator = GlobalAllocator(OnceLock::new());
    // try to randomly allocate and deallocate memory
    let mut rng = rand::thread_rng();
    for _ in 0..20 {
        let size = rng.gen_range(1..10000);
        let layout = std::alloc::Layout::from_size_align(size, 8).unwrap();
        let ptr = unsafe { global_allocator.alloc(layout) };
        unsafe { global_allocator.dealloc(ptr, layout) };
        println!("Allocated and deallocated {} bytes", size);
    }
    println!("Hello, world!");
}
